# _*_ coding:utf-8 _*_
# @File  : strategy.py
# @Time  : 2021-11-01  13:42
# @Author: zizle
import datetime
import pandas as pd
import db.utils as dbutils
from fastapi import APIRouter, Query, Path, Body
from db import FAConnection
from status import r_status
from category import VARIETY_NAME

admin_strategy_api = APIRouter()


def format_strategy_response(data_list):  # 格式化策略数据
    response_data = []
    for item in data_list:
        detail = eval(item['strategy'])

        # 转换数据类型
        sp, tp, cp = item['strategy_price'], item['target_price'], item['cutloss_price']
        sp = int(sp) if int(sp) == round(sp, 2) else round(sp, 2)
        tp = int(tp) if int(tp) == round(tp, 2) else round(tp, 2)
        cp = int(cp) if int(cp) == round(cp, 2) else round(cp, 2)

        # 添加策略状态
        # 设置策略状态 1未建仓(红色)，2运行中(绿色)，3已结束(灰色)
        state_text = '未建仓'
        if item['close_time']:
            state_text = '已结束'
        if item['open_time'] and not item['close_time']:
            state_text = '运行中'

        # 根据type显示策略文字内容
        if item['type'] == 1:  # 单边策略
            detail = detail[0]
            direction = '做多' if detail['direction'] == 'B' else '做空'
            strategy = f'【交易策略】建议于{sp}{detail["price_unit"]}附近{direction}{detail["contract"]}，' \
                       f'目标{tp}{detail["price_unit"]}，止损{cp}{detail["price_unit"]}。' \
                       f'{item["description"]}'
        elif item['type'] == 2:  # 双品种套利策略
            v1_data, v2_data = detail
            d1 = '做多' if v1_data["direction"] == 'B' else '做空'
            d2 = '做多' if v2_data['direction'] == 'B' else '做空'
            # 组织数据
            strategy = f'【套利策略】{v1_data["contract"]}与{v2_data["contract"]}合约价差为{sp}时，' \
                       f'建议{d1}{v1_data["contract"]},{d2}{v2_data["contract"]}，' \
                       f'价差为{tp}附近止盈，止损价差为{cp}附近。{item["description"]}'
        elif item['type'] == 3:  # 多品种套利策略
            strategy = '【套利策略】建议同时'
            for stv_item in detail:  # stv= strategy_variety
                d = '做多' if stv_item['direction'] == 'B' else '做空'
                strategy += f'{d}{stv_item["contract"]},数量:{stv_item["hands"]}手;'
            strategy = strategy[:-1] + '。'
            strategy += item['description']
        else:
            continue
        can_delete = True if (datetime.datetime.now() - item['create_time']).seconds <= 1800 else False
        open_time = item['open_time'].strftime('%Y-%m-%d %H:%M:%S') if item['open_time'] else ''
        close_time = item['close_time'].strftime('%Y-%m-%d %H:%M:%S') if item['close_time'] else ''

        profit = item['profit']
        profit = int(profit) if int(profit) == round(profit, 2) else round(profit, 2)
        response_data.append({
            'id': item['id'], 'type': item['type'], 'create_time': item['create_time'].strftime('%Y-%m-%d %H:%M:%S'),
            'creator': item['username'], 'open_time': open_time, 'close_time': close_time, 'profit': profit,
            'close_way': item['close_way'], 'can_delete': can_delete, 'content': strategy, 'state': state_text,
            'is_delete': item['is_delete']
        })
    return response_data


# 根据品种查询策略，品种可为空时查询所有策略,结果分页
@admin_strategy_api.get('/', summary='查询策略')
async def get_strategy(v: str = Query(0), page: int = Query(1, ge=1), page_size: int = Query(8, ge=5)):
    v = v if v else 0
    start, offset = (page - 1) * page_size, page_size
    sql = 'SELECT SQL_CALC_FOUND_ROWS a.id,a.create_time,a.type,a.variety_en,a.description,a.open_time,a.close_time,' \
          'a.close_way,a.profit,a.fund,a.strategy_price,a.target_price,a.cutloss_price,a.strategy,a.is_delete,' \
          'b.username ' \
          'FROM info_strategy AS a ' \
          'INNER JOIN user_user AS b ON a.creator=b.id ' \
          'WHERE IF("0"=%s,TRUE,a.variety_en=%s) ORDER BY a.create_time DESC LIMIT %s,%s;'
    db = FAConnection(conn_name='admin查询策略')
    records = db.query(sql, param=[v, v, start, offset], keep_conn=True)
    # 查询总条数
    total_obj = db.query('SELECT FOUND_ROWS() AS total;')[0]
    total_count = total_obj['total'] if total_obj else 1  # 当前数量
    # total_page = (total_count + page_size - 1) // page_size
    # 格式化数据返回
    response_data = format_strategy_response(records)
    return {'code': r_status.SUCCESS, 'message': '查询成功!', 'page': page, 'page_size': page_size,
            'total_count': total_count, 'data': response_data}


# 统计策略
@admin_strategy_api.get('/tj/', summary='统计策略')
async def statistics_strategy(v: str = Query(None)):
    def get_strategy_state(open_time, close_time):
        state = 1  # 未建仓
        if not pd.isna(close_time):
            state = 3  # 已结束
        if not pd.isna(open_time) and pd.isna(close_time):
            state = 2  # 运行中
        return state

    def tj_strategy(tj_df):
        data_obj = dict()
        data_obj['count'] = tj_df.shape[0]  # 数量
        data_obj['waiting'] = tj_df[tj_df['state'] == 1].shape[0]  # 等建仓
        data_obj['running'] = tj_df[tj_df['state'] == 2].shape[0]  # 运行中

        close_df = tj_df[tj_df['state'] == 3]  # 已结束
        data_obj['closed'] = close_df.shape[0]

        if data_obj['closed'] < 1:  # 没有已结束的策略
            data_obj['success_rate'] = ''
            data_obj['profit_rate'] = ''
        else:
            success_rate = close_df[close_df['profit'] > 0].shape[0] / data_obj['closed']
            data_obj['success_rate'] = str(round(success_rate * 100, 2)) + '%'
            profit_rate = close_df['profit'].sum() / close_df['fund'].sum()
            data_obj['profit_rate'] = str(round(profit_rate * 100, 2)) + '%'
        return data_obj

    v = v if v else 0
    sql = 'SELECT a.id,a.create_time,a.creator,a.type,a.variety_en,a.open_time,a.close_time,a.close_way,' \
          'a.profit,a.fund,a.strategy_price,a.target_price,a.cutloss_price,b.username,b.admin_name ' \
          'FROM info_strategy AS a ' \
          'INNER JOIN user_user AS b ON a.creator=b.id ' \
          'WHERE IF("0"=%s,TRUE,a.variety_en=%s) AND a.is_delete=0;'
    db = FAConnection(conn_name='统计策略查询')
    records = db.query(sql, param=[v, v])
    df = pd.DataFrame(records)
    if df.empty:
        return {'code': r_status.SUCCESS, 'message': '统计策略成功,没有查询到数据!',
                'year_data': [], 'month_data': [], 'variety_data': []}
    df['year'] = df['create_time'].apply(lambda x: x.year)
    df['month'] = df['create_time'].apply(lambda x: x.month)
    df['state'] = df.apply(lambda x: get_strategy_state(x['open_time'], x['close_time']), axis=1)
    year_groups = df.groupby(by=['year']).groups
    year_list = []
    # 按年统计(数量,未建仓,运行中,已结束,成功率,收益率)
    for y in year_groups:
        ydf = df[df['year'] == y]
        year_json = tj_strategy(ydf)
        year_json['year'] = y
        year_list.append(year_json)

    # --------- 按年统计数据结束
    # 按月统计(数量,未建仓,运行中,已结束,成功率,收益率)
    cur_ydf = df[df['year'] == datetime.datetime.today().year]
    month_text = ['一月', '二月', '三月', '四月', '五月', '六月', '七月', '八月', '九月', '十月', '十一月', '十二月']
    month_list = []
    for m in range(1, 13):
        mdf = cur_ydf[cur_ydf['month'] == m]
        month_json = tj_strategy(mdf)
        month_json['month'] = month_text[m - 1]
        month_list.append(month_json)

    # --------- 按月统计数据结束
    # 按品种统计
    variety_groups = df.groupby(by=['variety_en']).groups
    variety_list = []
    # 按品种统计(数量,未建仓,运行中,已结束,成功率,收益率)
    for v in variety_groups:
        vdf = df[df['variety_en'] == v]
        variety_json = tj_strategy(vdf)
        variety_json['variety_en'] = v
        variety_json['variety_name'] = VARIETY_NAME.get(v, v)
        variety_list.append(variety_json)

    # --------- 按人员统计数据结束
    username_list = df[['creator', 'admin_name']].to_dict(orient='records')
    username_dict = {uitem['creator']: uitem['admin_name'] for uitem in username_list}
    user_groups = df.groupby(by=['creator']).groups
    user_list = []
    # 按人员统计(数量,未建仓,运行中,已结束,成功率,收益率)
    for u in user_groups:
        udf = df[df['creator'] == u]
        user_json = tj_strategy(udf)
        user_json['creator'] = u
        user_json['username'] = username_dict.get(u, '未知')
        user_list.append(user_json)

    # 数量排序
    user_list.sort(key=lambda x: x['count'], reverse=True)
    return {'code': r_status.SUCCESS, 'message': '统计策略成功!',
            'year_data': year_list, 'month_data': month_list, 'variety_data': variety_list, 'user_data': user_list}


# 图形统计策略
@admin_strategy_api.get('/sy/', summary='策略收益sy')
async def strategy_profit(s: str = Query(None), e: str = Query(None), u: int = Query(None), v: str = Query(None)):
    # s - 开始日期, e-结束日期, u-人员id, v-品种
    try:
        if s:
            s = datetime.datetime.strptime(s, '%Y-%m-%d').strftime('%Y-%m-%d 00:00:00')
        if e:
            e = datetime.datetime.strptime(e, '%Y-%m-%d').strftime('%Y-%m-%d 00:00:00')
    except ValueError:
        return {'code': r_status.VALIDATE_ERROR, 'message': '参数错误', 'data': []}
    # 查询目标数据，根据日期处理每一日的累计收益率
    if not u:
        u = '0'
    if not v:
        v = '0'
    if not s:
        s = '0'
    if not e:
        e = '0'
    sql = 'SELECT id,close_time,profit,fund ' \
          'FROM info_strategy ' \
          'WHERE close_time IS NOT NULL AND ' \
          'IF("0"=%s,TRUE,close_time>=%s) AND IF("0"=%s,TRUE,close_time<=%s) ' \
          'AND IF("0"=%s,TRUE,creator=%s) AND IF("0"=%s,TRUE,variety_en=%s);'
    db_conn = FAConnection()
    records = db_conn.query(sql, param=[s, s, e, e, u, u, v, v])
    df = pd.DataFrame(records)
    if df.empty:
        return {'code': r_status.NOT_CONTENT, 'message': '没有数据', 'data': []}
    df['close_date'] = df['close_time'].apply(lambda x: x.strftime('%Y-%m-%d'))
    df['profit'] = df['profit'].astype(float)
    df['fund'] = df['fund'].astype(float)
    # 日期排序，求和
    sum_df = df.groupby(by=['close_date'], as_index=False).agg({'profit': 'sum', 'fund': 'sum'})
    sum_df.sort_values(by=['close_date'], inplace=True)
    # 滚动累计
    sum_df['cum_profit'] = sum_df['profit'].cumsum()
    sum_df['cum_fund'] = sum_df['fund'].cumsum()
    return {'code': r_status.SUCCESS, 'message': '数据图形统计成功', 'data': sum_df.to_dict(orient='records')}


# 根据ID查询单个策略详情数据
@admin_strategy_api.get('/{strategy_id}/', summary='ID查询策略')
async def get_strategy_by_id(strategy_id: int = Path(..., ge=1)):
    sql = 'SELECT * FROM info_strategy WHERE id=%s LIMIT 1;'
    db = FAConnection(conn_name='admin ID查询策略')
    obj = db.query(sql, param=[strategy_id], fetchone=True)[0]
    t = {1: '交易策略', 2: '2品种套利', 3: '多品种套利'}
    if obj:
        obj['create_time'] = obj['create_time'].strftime('%Y-%m-%d %H:%M:%S')
        obj['update_time'] = obj['update_time'].strftime('%Y-%m-%d %H:%M:%S')
        obj['type_text'] = t.get(obj['type'])
        return {'code': r_status.SUCCESS, 'message': '查询成功！', 'data': obj}
    return {'code': r_status.NOT_CONTENT, 'message': '数据不存在!', 'data': {}}


# 根据ID修改单个策略的详情数据
@admin_strategy_api.put('/{strategy_id}/', summary='ID修改策略')
async def modify_strategy_by_id(strategy_id: int = Path(..., ge=1), body_item: dict = Body(...)):
    # 去掉None
    body_item = {k: body_item[k] for k in filter(lambda x: body_item[x] is not None, body_item)}
    accept_columns = ['fund', 'profit', 'open_time', 'close_time', 'close_way', 'strategy_price', 'target_price',
                      'cutloss_price', 'strategy', 'is_open', 'is_delete']
    s_item = dbutils.exclude_columns(body_item, accept_columns)
    # 生成sql
    middle_sql = dbutils.create_sql(s_item)
    sql = 'UPDATE info_strategy SET ' + middle_sql + 'WHERE id=%(id)s LIMIT 1;'
    s_item['id'] = strategy_id
    db = FAConnection(conn_name='根据ID修改策略')
    _, s = db.execute(sql, param=s_item)
    if s:
        return {'code': r_status.SUCCESS, 'message': '修改数据成功!'}
    return {'code': r_status.UNKNOWN_ERROR, 'message': '修改数据失败,原因未知!'}
